#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
知识搜索向量恢复修复脚本

问题分析：
1. ChromaDB的knowledge集合为空，但documents集合有5个文档
2. knowledge_search使用RAG管道，依赖vectorizer的内存向量存储
3. 当内存向量为空时，应该从ChromaDB的documents集合恢复向量
4. 但恢复过程可能存在问题

修复方案：
1. 检查ChromaDB中documents集合的数据结构
2. 手动触发向量恢复过程
3. 验证向量恢复是否成功
4. 测试知识搜索功能
"""

import asyncio
import sys
import os
import json
from datetime import datetime
from pathlib import Path

# 添加项目根目录到Python路径
project_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root))

from src.core.context_orchestrator import ContextOrchestrator
from src.config import load_config

async def fix_knowledge_search_vectors():
    """修复知识搜索向量恢复问题"""
    
    config = load_config()
    orchestrator = ContextOrchestrator(config)
    
    try:
        await orchestrator.initialize()
        
        print("=== 知识搜索向量恢复修复 ===")
        
        # 1. 检查ChromaDB documents集合的详细信息
        print("\n1. 检查ChromaDB documents集合详细信息")
        chromadb_client = orchestrator.db_manager.chromadb_client
        
        try:
            collection = chromadb_client.get_collection('documents')
            if collection:
                # 获取集合中的所有数据
                all_data = collection.get(include=['embeddings', 'metadatas', 'documents'])
                
                print(f"documents集合中的文档数量: {len(all_data['ids'])}")
                print(f"包含向量的文档数量: {sum(1 for emb in all_data['embeddings'] if emb is not None)}")
                
                # 显示前几个文档的信息
                for i, (doc_id, metadata, document) in enumerate(zip(
                    all_data['ids'][:3], all_data['metadatas'][:3], all_data['documents'][:3]
                )):
                    print(f"\n文档 {i+1}:")
                    print(f"  ID: {doc_id}")
                    print(f"  元数据: {metadata}")
                    print(f"  内容预览: {document[:100]}...")
                    print(f"  有向量: {all_data['embeddings'][i] is not None}")
            else:
                print("documents集合不存在")
        except Exception as e:
            print(f"检查documents集合失败: {e}")
        
        # 2. 检查knowledge_processor的vectorizer状态
        print("\n2. 检查knowledge_processor的vectorizer状态")
        knowledge_processor = orchestrator.knowledge_processor
        
        if knowledge_processor.vectorizer:
            vector_store = knowledge_processor.vectorizer.vector_store
            if hasattr(vector_store, 'vectors'):
                print(f"内存向量存储中的向量数量: {len(vector_store.vectors)}")
                print(f"向量存储类型: {type(vector_store)}")
            else:
                print("向量存储没有vectors属性")
        else:
            print("vectorizer未初始化")
        
        # 3. 手动触发向量恢复
        print("\n3. 手动触发向量恢复")
        try:
            await knowledge_processor._restore_vectors_from_chromadb()
            
            # 检查恢复后的状态
            if knowledge_processor.vectorizer and hasattr(knowledge_processor.vectorizer.vector_store, 'vectors'):
                vector_count = len(knowledge_processor.vectorizer.vector_store.vectors)
                print(f"向量恢复后，内存中的向量数量: {vector_count}")
            else:
                print("向量恢复后，仍无法访问向量存储")
                
        except Exception as e:
            print(f"向量恢复失败: {e}")
            import traceback
            traceback.print_exc()
        
        # 4. 测试知识搜索功能
        print("\n4. 测试知识搜索功能")
        
        test_queries = [
            "Python编程语言",
            "Vue.js框架",
            "编程",
            "框架"
        ]
        
        search_results = {}
        
        for query in test_queries:
            try:
                print(f"\n测试查询: {query}")
                
                # 使用RAG管道搜索
                rag_results = await knowledge_processor.search_knowledge(
                    query, user_id="default_user", limit=3, use_rag_pipeline=True
                )
                print(f"  RAG管道结果数量: {len(rag_results)}")
                
                # 使用传统搜索
                legacy_results = await knowledge_processor.search_knowledge(
                    query, user_id="default_user", limit=3, use_rag_pipeline=False
                )
                print(f"  传统搜索结果数量: {len(legacy_results)}")
                
                search_results[query] = {
                    'rag_count': len(rag_results),
                    'legacy_count': len(legacy_results),
                    'rag_results': rag_results[:1] if rag_results else [],  # 只保存第一个结果
                    'legacy_results': legacy_results[:1] if legacy_results else []
                }
                
            except Exception as e:
                print(f"  搜索失败: {e}")
                search_results[query] = {'error': str(e)}
        
        # 5. 生成修复报告
        print("\n5. 生成修复报告")
        
        report = {
            'timestamp': datetime.now().isoformat(),
            'chromadb_documents_count': len(all_data['ids']) if 'all_data' in locals() else 0,
            'vectorizer_initialized': knowledge_processor.vectorizer is not None,
            'memory_vectors_count': len(knowledge_processor.vectorizer.vector_store.vectors) if (
                knowledge_processor.vectorizer and 
                hasattr(knowledge_processor.vectorizer.vector_store, 'vectors')
            ) else 0,
            'search_test_results': search_results,
            'fix_status': 'completed'
        }
        
        # 保存报告
        with open('knowledge_search_fix_report.json', 'w', encoding='utf-8') as f:
            json.dump(report, f, ensure_ascii=False, indent=2)
        
        print("修复报告已保存到 knowledge_search_fix_report.json")
        
        # 6. 提供修复建议
        print("\n=== 修复建议 ===")
        
        if report['memory_vectors_count'] > 0:
            print("✅ 向量恢复成功，知识搜索应该可以正常工作")
        else:
            print("❌ 向量恢复失败，需要进一步调查")
            print("建议检查：")
            print("1. ChromaDB中的向量数据是否完整")
            print("2. vectorizer的初始化是否正确")
            print("3. 向量恢复逻辑是否有bug")
        
        # 检查搜索结果
        successful_searches = sum(1 for result in search_results.values() 
                                if 'error' not in result and result.get('rag_count', 0) > 0)
        
        if successful_searches > 0:
            print(f"✅ {successful_searches}/{len(test_queries)} 个测试查询返回了结果")
        else:
            print("❌ 所有测试查询都没有返回结果")
        
    finally:
        await orchestrator.shutdown()

if __name__ == "__main__":
    asyncio.run(fix_knowledge_search_vectors())